package org.shj.spark.streaming

import org.apache.spark.SparkConf
import org.apache.spark.SparkContext
import org.apache.spark.streaming.Seconds
import org.apache.spark.streaming.State
import org.apache.spark.streaming.StreamingContext
import org.apache.spark.streaming.StateSpec

/**
 * 在流式处理中，时间间隔后的任务只处理此时间间隔内收到的数据。如果想要在此次任务中合并上一次处理的结果，
 * 我们可以用 updateStateByKey 这个 transformation。此时，需要设置 checkpoint
 * 
 * 此例演示了两种获取上一次处理结果，并把结果合并到本次任务的方式：
 * 方式一：   
 * ./bin/spark-submit --master spark://master:7077 --class org.shj.spark.streaming.UpdateStateByKeyDemo /opt/tmp/SparkDemo.jar
 *  1 /spark/streaming /spark/checkpoint 
 * 
 * 方式二：
 * ./bin/spark-submit --master spark://master:7077 --class org.shj.spark.streaming.UpdateStateByKeyDemo /opt/tmp/SparkDemo.jar 
 * 2 /spark/checkpoint 192.168.56.101 9999
 */
object UpdateStateByKeyDemo {
  def main(args: Array[String]): Unit = {
    if (args.length < 3) {
      System.err.println("Usage: UpdateStateByKeyDemo <method> <directory> <checkpointDirectory>")
      System.exit(1)
    }
    
    val conf = new SparkConf().setAppName("UpdateStateByKeyDemo")
    //conf.setMaster("local[*]")
    val sc = new SparkContext(conf)
    sc.setLogLevel("WARN")
    
    if(args(0) == "1")
      updateStateMethod1(sc, args(1), args(2))
    else
      updateStateMethod2(sc, args(1), args(2), args(3).toInt)
    
  }
  
  def updateStateMethod1(sc: SparkContext, processDir: String, checkPointDir: String){
    val ssc = new StreamingContext(sc, Seconds(10)) 
    ssc.checkpoint(checkPointDir) //下面的 updateStateByKey 需要设置 checkpoint
    
    val lines = ssc.textFileStream(processDir)
    
    val pairs = lines.flatMap(_.split("\\s+")).map(word => (word, 1))
    
    //方法不能做为参数进行传递，而函数则可以，所以下面的下划线("_") 的作用是把 updateFunc 这个方法转化成函数
    val runningCount = pairs.updateStateByKey(updateFunc _) 
    
    val wordCnt = runningCount.reduceByKey(_+_)
    wordCnt.print()
    
    /**
     * 此处的设置从注释来看，以为是每次处理后把结果存在以指定前辍和后辍的文件中。其实这里指的是目录。目录名的格式是 prefix-TIME_IN_MS.suffix
     * 此目录是生成在HDFS上的/user/root 目录(用的是root的帐号提交的程序)下
     * TODO: 如何指定这里目录是存在哪个目录下？
     */
    wordCnt.saveAsTextFiles("ssc", "txt")
    
    ssc.start()
    ssc.awaitTermination()
  }
  
  /**
   * 此方式中使用 mapWithState, 从其注释中可以得出：
   * 此方式下，只会从上次时间间隔内取出具有和此次任务相同的key的值进行mapFunction的操作。即如果此次任务中，有3个key： I, love, you; 
   * 而上一次的统计结果是： (I, 2), (hate, 1), (you 2), 那么此次的运行结果是(I, 3), (love, 1), (you 3)。 也就是说(hate,1) 将会被抛弃
   */
  def updateStateMethod2(sc: SparkContext, checkPointDir: String, ip: String, port: Int) = {
    val ssc = new StreamingContext(sc, Seconds(10)) 
    ssc.checkpoint(checkPointDir)
    
    val lines = ssc.socketTextStream(ip, port)
    val words = lines.flatMap(_.split("\\s+"))
    val wordDstream = words.map(x => (x, 1))
    
    val mapFunc = (word: String, one: Option[Int], state: State[Int]) => {
      val sum = one.getOrElse(0) + state.getOption().getOrElse(0)
      val output = (word, sum)
      state.update(sum)
      output
    }
    
    val stateDstream = wordDstream.mapWithState(StateSpec.function(mapFunc))
    stateDstream.print()
    
    ssc.start()
    ssc.awaitTermination()
  }
  
  /**
   * updateStateByKey 中所需要的Function， 此Function中的第一个参数是此次时间间隔内的value的值的结果的集合，
   * 第二个参数是上一次时间间隔内的value的值的结果。
   * 
   */
  def updateFunc(newValues: Seq[Int], previousCount: Option[Int]): Option[Int] = {
    val newCount = newValues.sum + previousCount.getOrElse(0)
    Some(newCount)
  }
}